package com.tuniu.agents.common.logger;

import com.tuniu.agents.advisor.AgentChatMemoryAdvisor;
import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.common.IdGenerator;
import com.tuniu.agents.common.TraceEnum;
import com.tuniu.agents.common.TraceLogVo;
import com.tuniu.agents.common.enums.SqlOperationType;
import com.tuniu.agents.message.TraceLogger;
import com.tuniu.agents.util.JsonUtil;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.ProceedingJoinPoint;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import java.util.Map;

/**
 * Records request information for all tools
 */
@Aspect
@Component
@Order(2)
public class TraceDateAspect {

    @Autowired
    private TraceLogger traceLogger;

    @Autowired
    private IdGenerator idGenerator;

    @Autowired(required = false)
    private ChatEnvironmentMemory chatEnvironmentMemory;

    private static final Logger logger = LoggerFactory.getLogger(TraceDateAspect.class);

    @Around("execution(* com.tuniu.agents.tools..*.apply(..))")
    public Object aroundApply(ProceedingJoinPoint joinPoint) throws Throwable {
        // Get tool information
        Class<?> targetClass = joinPoint.getTarget().getClass();
        String toolName = "";
        try {
            toolName = (String) targetClass.getField("TOOLNAME").get(null);
        } catch (Exception e) {
            logger.warn("Field TOOLNAME not found in class: {}", targetClass.getName());
        }

        // Get method signature and method object
        Object[] args = joinPoint.getArgs();
        String conversationId = getConversationId(args);
        saveTrace(conversationId, toolName, args);
        // Print all input parameters
        return joinPoint.proceed(args);
    }

    private void saveTrace(String conversationId, String toolName, Object[] args) {
        try {
            Map<String, Object> toolContext = JsonUtil.toMap(JsonUtil.toJson(args[1]));
            Object agentId = JsonUtil.toMap(JsonUtil.toJson(toolContext.get("context"))).get(AgentChatMemoryAdvisor.CHAT_MEMORY_AGENT_ID_KEY);
            Map<String, Object> map = chatEnvironmentMemory.get(conversationId);
            TraceLogVo traceLogVo = new TraceLogVo.Builder()
                    .traceId(map.getOrDefault(AgentsConstants.AgentsConversationEnvironmentKey.PARAM_TRACE_ID,"").toString())
                    .userId(map.get(AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID) != null ? map.get(AgentsConstants.AgentsConversationEnvironmentKey.CUSTOMER_ID).toString() : null)
                    .conversationId(conversationId)
                    .agentId(agentId != null ? agentId.toString() : "")
                    .tool(toolName)
                    .type(TraceEnum.TOOL_INPUT.getDescription())
                    .content(JsonUtil.toJson(args[0]))
                    .toolContext(JsonUtil.toJson(args[1]))
                    .operationType(SqlOperationType.ChatTraceLogSaveTrace.getDescription())
                    .build();
            traceLogger.saveTrace(traceLogVo);
        } catch (Exception e) {
            logger.error("TraceDateAspect.saveTrace error", e);
        }
    }

    public String getConversationId(Object[] args) {
        Object conversationId = ((ToolContext) args[1]).getContext().get(AgentsConstants.AgentsConversationEnvironmentKey.CONVERSATION_ID);
        return conversationId == null ? null : String.valueOf(conversationId);
    }

}
